#include <stdlib.h>
#include <stdio.h>
#include <pthread.h>
#include <sys/time.h>

#define a(i,j) a[(i)*k + (j)]
#define b(i,j) b[(i)*n + (j)]
#define c(i,j) c[(i)*n + (j)]

// packing

void pack_A(float *a, int k, int m, float *pakage_a){
    float *tep = &a(m,0);
    for (int i = 0; i < k; i++){
        pakage_a[i] = tep[i];
        pakage_a[i+k] = tep[i+k];
    }
}

void pack_B(float *b, int k, int n, int col, float *pakage_b){
   for (int i = 0; i < k; i++){
        int s = i*n + col;
        pakage_b[i] = b[s];
        pakage_b[i+k] = b[s +1];
   }
}

void matmul_2_2_kernel(float *pakage_a, float *pakage_b, float *c, int x, int y, int k, int n){
    register float c00,c01,c10,c11;
    c00 = 0.0;
    c01  = 0.0;
    c10  = 0.0;
    c11  = 0.0;
    for (int s = 0; s < k; s++){
        int a1 = k + s;
        c00 += pakage_a[s] * pakage_b[s];
        c01 += pakage_a[s] * pakage_b[a1];
        c10 += pakage_a[a1] * pakage_b[s];
        c11 += pakage_a[a1] * pakage_b[a1];
    }
    
    c(x,y) = c00;
    c(x,y+1) = c01;
    c(x+1,y) = c10;
    c(x+1,y+1) = c11;
}

// gemm
void matmul_2_2(int m, int n, int k, float *a, float *b, float *c){
    float *p_a = (float *)malloc(2*k*sizeof(float));
    float *p_b = (float *)malloc(2*k*sizeof(float));
    for (int i = 0; i < m; i+=2){
        pack_A(a, k, i, p_a);
        for (int j = 0; j < n; j+=2){
            pack_B(b, k, n, j, p_b);
            matmul_2_2_kernel(p_a, p_b, c, i, j, k, n);
        }
    }
    
    free(p_a);
    free(p_b);
}

// multithread

typedef struct Mes{
    float *a;
    float *b;
    float *c;
    int m;
    int n;
    int k;
}*mes_ptr;

void *matmul_2_2_thread_function(void *arg){
    mes_ptr mes = (struct Mes *)arg;
    matmul_4_4(mes -> m, mes -> n, mes -> k, mes -> a, mes -> b, mes -> c);

}

void mul_4_4_mp(int m, int n, int k, float *a, float *b, float *c){
    const int num_threads = 8;
    struct Mes mes[num_threads];
    pthread_t threads[num_threads];
    int m_nums = m/num_threads;
    
    for (int i = 0; i < num_threads; i++){
        mes[i].a = &a(i*m_nums,0);
        mes[i].b = b;
        mes[i].c = &c(i*m_nums,0);
        mes[i].m = m_nums;
        mes[i].n = n;
        mes[i].k = k;
    }

    for(int i = 0; i < num_threads; i++){
        pthread_create(&threads[i], NULL, matmul_2_2_thread_function, &mes[i]);
    }

    for (int i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }

}

int main(){

    int m = 2048, n = 2048, k = 2048;
    float *a = (float *)malloc(m*k*sizeof(float));
    float *b = (float *)malloc(k*n*sizeof(float));
    float *c = (float *)malloc(m*n*sizeof(float));

    for (int i = 0; i < m*k; i++){
        a[i] = 1.0f;
    }

    for (int i = 0; i < k*n; i++){
        b[i] = 1.0f;
    }

    for (int i = 0; i < m*n; i++){
        c[i] = 0.0f;
    }

    struct timeval start, end;
    long seconds_finished;
    unsigned long micro_seconds_finished;
    double elapsed;
    gettimeofday(&start, NULL);

    // function


    gettimeofday(&end, NULL);
    seconds_finished = end.tv_sec - start.tv_sec;
    micro_seconds_finished = end.tv_usec - start.tv_usec;
    elapsed = ((seconds_finished * 1000000) + micro_seconds_finished) / 1000000.0;
    printf("time: %f s\n", elapsed);

    double flops = 2.0 * m * n * k * 1e-9;
    printf("flops: %lf Gflops\n", flops/elapsed);

    free(a);
    free(b);
    free(c);

    return 0;
}